import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from data_set import stock_prediction_data_set
from su_she_model import TransformerClassifier

# 检测是否支持 CUDA
if torch.cuda.is_available():
    print("CUDA is available.")
    # 这里可以进行使用 GPU 的相关操作
    device = torch.device("cuda")
else:
    print("CUDA is not available. Using CPU.")
    device = torch.device("cpu")

window_size = 30
# 创建自定义数据集实例
custom_dataset = stock_prediction_data_set(
    data_dir="data",
    data_file_limit=0.01,
    window=window_size + 1,
    data_column=['openIndex', 'closeIndex', 'highIndex', 'lowIndex',
                 'turnoverValueIndex', 'turnoverVolIndex',
                 'z2', 'z10', 'z30', 'm2', 'm10', 'm30',
                 'openHl', 'closeHl', 'highHl', 'lowHl'],
    label_column="profit")

# 创建 DataLoader
batch_size = 30

# 划分训练集和验证集
validate_loader = DataLoader(custom_dataset, batch_size=batch_size, shuffle=True)

# 模型和参数
num_layers = 10

d_model = 16
num_head = 8
dim_feedforward = 512
ff_hidden_dim = 128
output_dim = 2
dropout = 0.1

# 创建模型
model = TransformerClassifier(num_layers, d_model, window_size, num_head, dim_feedforward, ff_hidden_dim,
                              output_dim).to(
    device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
# 加载已经训练好的模型
model.load_state_dict(torch.load('transformer_classifier-80%.pth', map_location='cpu'))
# 在验证集上评估模型
model.eval()  # 设置模型为评估模式
val_loss = 0
correct = 0
total = 0

with torch.no_grad():
    for val_batch_data, val_batch_labels in validate_loader:
        val_batch_labels = val_batch_labels.to(device)
        val_outputs = model(val_batch_data.to(device))
        val_loss += criterion(val_outputs, val_batch_labels).item()
        predicted = torch.where(val_outputs > 0.5, torch.tensor(1.0), torch.tensor(0.0))
        total += val_batch_labels.size(0)
        correct += ((predicted == val_batch_labels).sum().item() / 2)

# 打印在验证集上的性能指标
val_accuracy = correct / total
val_loss_result = val_loss / len(validate_loader)
print(f"Validation Loss: {val_loss_result}, Validation Accuracy: {val_accuracy}")
